import csv
import os

import requests
from lxml import etree

from utils.random_request import random_ua


def init():
    """
    初始化数据文件，如果没有文件夹则创建，并写入表头
    注意：如果文件已经存在，则会覆盖原文件
    :return:
    """
    if not os.path.exists("./data"):
        os.mkdir("./data")
    with open("./data/city_url_houseNum.csv", "w", newline="", encoding="utf-8") as f:
        writer = csv.writer(f)
        writer.writerow(["city_name", "city_url", "house_num"])


def writeRow(city_name, city_url, house_num):
    """
    将数据写入csv文件
    :param city_name:
    :param city_url:
    :param house_num:
    :return:
    """
    with open("./data/city_url_houseNum.csv", "a", newline="", encoding="utf-8") as f:
        writer = csv.writer(f)
        writer.writerow([city_name, city_url, house_num])


def read_city_url_from_csv():
    """
    读取城市url列表
    :return:
    """
    with open("./data/city_url.csv", "r", encoding="utf-8") as f:
        reader = csv.reader(f)
        # 跳过文件头部
        next(reader)
        for row in reader:
            yield row


def get_house_num(url):
    """
    获取城市对应的房源数量
    :param url:
    :return:
    """
    headers = {"User-Agent": random_ua()}
    try:
        response = requests.get(url, headers=headers)
        if response.status_code == 200:
            html = response.text
            root = etree.HTML(html)
            house_num = root.xpath("/html/body/div[3]/div[2]/div/span[2]/text()")[0]
            return house_num

    except requests.exceptions.RequestException as e:
        return None


def main():
    for city_name, city_url in read_city_url_from_csv():
        house_num = get_house_num(city_url + "/loupan/")
        if house_num is not None and house_num.isdigit():
            writeRow(city_name, city_url, house_num)


if __name__ == "__main__":
    init()
    main()
